##########################################################################################

library(dplyr)
library(ggplot2)
library(data.table)
library(RColorBrewer)
library(optparse)

##########################################################################################

option_list <- list(
    make_option(c("--maf_cancer_file"), type = "character") ,
    make_option(c("--maf_im_file"), type = "character") ,
    make_option(c("--images_path"), type = "character") ,
    make_option(c("--info_file"), type = "character")
)

if(1!=1){


}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

maf_cancer_file <- opt$maf_cancer_file
maf_im_file <- opt$maf_im_file
info_file <- opt$info_file
images_path <- opt$images_path

dir.create(images_path , recursive = T)

###########################################################################################

info <- data.frame(fread(info_file))
dat_maf_cancer <- data.frame(fread( maf_cancer_file ))
dat_maf_im <- data.frame(fread( maf_im_file ))

###########################################################################################

Variant_Type <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")

###########################################################################################
## 不看IGC + DGC 均有的样本
## 只看MSS的样本
info_gc_divide <- subset( info , Class != "IGC + DGC" & MS_Type == "MSS" ) 
info_gc <- info_gc_divide
info_gc$Class <- "GC"
info_im <- subset( info , Class != "IGC + DGC" & From == "NJMU" & MS_Type == "MSS" ) 
info_im$Class <- "IM"

info_use <- rbind( info_gc_divide , info_gc , info_im )
info_all <- info_use
info_all$From <- 'All'
info_use <- rbind(info_all , info_use)

###########################################################################################
## 注释肿瘤样本的类型
maf <- dat_maf_cancer
maf_use <- data.frame(
    Hugo_Symbol = maf$Hugo_Symbol, 
    Chromosome = maf$Chromosome , Start_Position =  maf$Start_position , 
    Reference_Allele = maf$Reference_Allele , Tumor_Seq_Allele2 = maf$Tumor_Seq_Allele2 , 
    Variant_Classification = maf$Variant_Classification , Tumor = maf$Tumor_Sample_Barcode , From = maf$From )

maf_use <- subset( maf_use , Variant_Classification %in% Variant_Type )
maf_use <- merge( maf_use , info[,c("Tumor" , "Class")] , by = "Tumor")
maf_cancer <- maf_use

###########################################################################################
## 肠化样本
maf <- dat_maf_im
maf_use <- data.frame(
    Hugo_Symbol = maf$Hugo_Symbol, 
    Chromosome = maf$Chromosome , Start_Position =  maf$Start_position , 
    Reference_Allele = maf$Reference_Allele , Tumor_Seq_Allele2 = maf$Tumor_Seq_Allele2 , 
    Variant_Classification = maf$Variant_Classification , Tumor = maf$Tumor_Sample_Barcode , From = maf$From )

maf_use <- subset( maf_use , Variant_Classification %in% Variant_Type )
maf_im <- maf_use
maf_im$Class <- "IM"
maf_im <- subset( maf_im , Tumor %in% info_im$Tumor  )

###########################################################################################
## 总的，不同来源的
## IGC和DGC分开的，合并的
maf_cancer2 <- maf_cancer
maf_cancer2$Class <- "GC"

maf_use <- rbind(maf_cancer , maf_cancer2 , maf_im)
maf_use$From <- "All"
maf_use <- rbind(maf_use , maf_cancer , maf_cancer2 , maf_im)

maf_use <- subset( maf_use , Tumor %in% info_use$Tumor  )
maf_use <- merge(maf_use,info_use[,c("Tumor","Age","Tobacco","Alcohol","HP","Molecular.subtype")],by="Tumor")

###########################################################################################
## 按照病理类型计算突变率
## 以人为单位
mutRate_use_baseline <- function(maf_use,info_use,use_baseline){
    mut_rate <- maf_use %>%
    group_by( Hugo_Symbol , Class , From,!!sym(use_baseline) ) %>%
    summarize( MutNum = length(unique(Tumor)) )

    mut_rate$id <- paste0( mut_rate$Class , ":" , mut_rate$From,":",mut_rate[[use_baseline]] )

    class_num <- info_use %>%
    group_by( Class , From,!!sym(use_baseline)) %>%
    summarize( SampleNum = length(unique(Tumor)) )
    class_num$id <- paste0( class_num$Class , ":" , class_num$From ,":",class_num[[use_baseline]])

    mut_rate <- merge( mut_rate , class_num[,c("id" , "SampleNum")] , by = "id" )
    ## 分IGC、DGC和IM
    mut_rate$MutRate <- mut_rate$MutNum/mut_rate$SampleNum

    return(mut_rate)
}

mutRate_use_baseline_subtype <- function(maf_use,info_use,use_baseline){
    mut_rate <- maf_use %>%
    group_by( Hugo_Symbol , Class , From,!!sym(use_baseline), Molecular.subtype ) %>%
    summarize( MutNum = length(unique(Tumor)) )

    mut_rate$id <- paste0( mut_rate$Class , ":" , mut_rate$From,":",mut_rate[[use_baseline]],":",mut_rate$Molecular.subtype )

    class_num <- info_use %>%
    group_by( Class , From,!!sym(use_baseline),Molecular.subtype) %>%
    summarize( SampleNum = length(unique(Tumor)) )
    class_num$id <- paste0( class_num$Class , ":" , class_num$From ,":",class_num[[use_baseline]],":",class_num$Molecular.subtype)

    mut_rate <- merge( mut_rate , class_num[,c("id" , "SampleNum")] , by = "id" )
    ## 分IGC、DGC和IM
    mut_rate$MutRate <- mut_rate$MutNum/mut_rate$SampleNum

    return(mut_rate)
}

##吸烟
#全部的
mut_rate_Tobacco <- mutRate_use_baseline(maf_use,info_use,use_baseline="Tobacco")
mut_rate_Tobacco$Molecular.subtype <- "MSS"
#CIN和GS的
mut_rate_Tobacco_subtype <- mutRate_use_baseline_subtype(maf_use,info_use,use_baseline="Tobacco")

mut_rate_Tobacco_all <- rbind(mut_rate_Tobacco,mut_rate_Tobacco_subtype)
out_name <- paste0( images_path , "/MutRate_baline_Tobacco.tsv" )
write.table( mut_rate_Tobacco_all , out_name , row.names = F , quote = F , sep = "\t" )


##饮酒
#全部的
mut_rate_Alcohol <- mutRate_use_baseline(maf_use,info_use,use_baseline="Alcohol")
mut_rate_Alcohol$Molecular.subtype <- "MSS"

mut_rate_Alcohol_subtype <- mutRate_use_baseline_subtype(maf_use,info_use,use_baseline="Alcohol")
mut_rate_Alcohol_all <- rbind(mut_rate_Alcohol,mut_rate_Alcohol_subtype)


out_name <- paste0( images_path , "/MutRate_baline_Alcohol.tsv" )
write.table( mut_rate_Alcohol_all , out_name , row.names = F , quote = F , sep = "\t" )


##hp

mut_rate_HP <- mutRate_use_baseline(maf_use,info_use,use_baseline="HP")
mut_rate_HP$Molecular.subtype <- "MSS"

mut_rate_HP_subtype <- mutRate_use_baseline_subtype(maf_use,info_use,use_baseline="HP")
mut_rate_HP_all <- rbind(mut_rate_HP,mut_rate_HP_subtype)
out_name <- paste0( images_path , "/MutRate_baline_HP.tsv" )
write.table( mut_rate_HP_all , out_name , row.names = F , quote = F , sep = "\t" )

##年龄

mut_rate_HP <- mutRate_use_baseline(maf_use,info_use,use_baseline="Age")
mut_rate_HP$Molecular.subtype <- "MSS"

mut_rate_HP_subtype <- mutRate_use_baseline_subtype(maf_use,info_use,use_baseline="Age")
mut_rate_HP_all <- rbind(mut_rate_HP,mut_rate_HP_subtype)
out_name <- paste0( images_path , "/MutRate_baline_Age.tsv" )
write.table( mut_rate_HP_all , out_name , row.names = F , quote = F , sep = "\t" )
